'''
author:        Wang Chenyang <cy-wang21@mails.tsinghua.edu.cn>
date:          2025-07-17
Copyright © Department of Physics, Tsinghua University. All rights reserved
'''

import unreal
import numpy as np

# Not used
UE5_SCALING = 100 # 现实中的1m对应UE5中的多少cm
R_EARTH = 6371e3 # 地球半径
COORD_TRANS = np.eye(2) # 坐标变换. 若UE5中x轴指向正东，y轴指向正北，则无需修改.
NAME_DATA_MAP = {}
SKIP_ROWS = 1   # CSV文件中需要跳过几行
SEQ_DIR = ""    # Sequences的存储路径
MASTER_SEQUENCE_NAME = ""
ROLL_REVERSE = False
PITCH_REVERSE = False
YAW_REVERSE = False

def global_to_local(
    longitudes: dict[np.ndarray[1, float]],
    latitudes: dict[np.ndarray[1, float]]
):
    '''
    Args:
        longitudes: 所有飞机的经度
        latitudes: 所有飞机的纬度 
    Return:
        (phi_bar, theta_bar): 平均经度与平均纬度
        X: x坐标(以正东方向为正)
        Y: y坐标(以正北方向为正)
        所有长度坐标以实际m为单位
    '''

    # 1. 求平均经纬度
    phi_bar = 0
    theta_bar = 0
    for plane_idx in longitudes.keys():
        phi_bar += np.mean(longitudes[plane_idx])
        theta_bar += np.mean(latitudes[plane_idx])
    phi_bar /= len(longitudes)
    theta_bar /= len(latitudes)

    # 2. 求x, y坐标
    X = {}
    Y = {}
    for plane_idx in longitudes.keys():
        rel_log = longitudes[plane_idx] - phi_bar
        rel_lat = latitudes[plane_idx] - theta_bar
        curr_X = R_EARTH * np.cos(theta_bar) * rel_log
        curr_Y = R_EARTH * rel_lat
        # 3. 旋转
        X[plane_idx] = (
            COORD_TRANS[0,0] * curr_X + COORD_TRANS[0,1] * curr_Y
        )
        Y[plane_idx] = (
            COORD_TRANS[1,0] * curr_X + COORD_TRANS[1,1] * curr_Y
        )
    
    return (phi_bar, theta_bar), X, Y


def load_DCS_data():
    '''
    Args:
        DCS_data_path: DCS数据路径
    Return:
        获得所有飞机的
        (时间, 经度, 纬度, 高度, Roll, Pitch, Yaw)
    '''
    all_time = {}
    all_longitudes = {}
    all_latitudes = {}
    all_altitudes = {}
    all_rolls = {}
    all_pitches = {}
    all_yaws = {}

    for plane_idx in NAME_DATA_MAP.keys():
        data = np.loadtxt(
            NAME_DATA_MAP[plane_idx],
            delimiter=',',
            skiprows=SKIP_ROWS,
            dtype=np.float64
        )
        all_time[plane_idx] = data[:,0]
        all_longitudes[plane_idx] = data[:,1] * np.pi / 180 # to rad
        all_latitudes[plane_idx] = data[:,2] * np.pi / 180 # to rad
        all_altitudes[plane_idx] = data[:,3]
        all_rolls[plane_idx] = data[:,4]
        all_pitches[plane_idx] = data[:,5]
        all_yaws[plane_idx] = data[:,6]

        if ROLL_REVERSE:
            all_rolls[plane_idx] = -all_rolls[plane_idx]
        if PITCH_REVERSE:
            all_pitches[plane_idx] = -all_pitches[plane_idx]
        if YAW_REVERSE:
            all_yaws[plane_idx] = -all_yaws[plane_idx]

    return (
        all_time,
        all_longitudes,
        all_latitudes,
        all_altitudes,
        all_rolls,
        all_pitches,
        all_yaws
    )

# MODIFIED: The function now accepts the master level sequence as an argument
def add_actor_to_sequence(level_sequence, act_name, t, X, Y, Z, Roll, Pitch, Yaw):
    '''
    Args:
        level_sequence: The master Level Sequence to add the animation to
        act_name: 飞机名称
        t: 时间
        X: x坐标
        Y: y坐标
        Z: z坐标
        Roll: 横滚角度
        Pitch: 俯仰角度
        Yaw: 偏航角度
    Return:
        None
    '''
    editor_world = unreal.get_editor_subsystem(unreal.UnrealEditorSubsystem).get_editor_world()
    actor_label = act_name # 【请确保关卡中Actor的标签与此一致】
    actor_to_animate = None
    all_actors = unreal.GameplayStatics.get_all_actors_of_class(editor_world, unreal.Actor)
    for actor in all_actors:
        if actor.get_actor_label() == actor_label:
            actor_to_animate = actor
            break

    if not actor_to_animate:
        unreal.log_error(f"错误：在关卡中找不到名为 '{actor_label}' 的Actor。")
        return # MODIFIED: Exit the function if actor not found
    else:
        unreal.log(f"成功找到Actor: {actor_to_animate.get_actor_label()}，并将其添加到序列中。")

    # -----------------------------------------------------------------------------
    # 1. 将Actor绑定到序列
    # -----------------------------------------------------------------------------
    # MODIFIED: We now use the provided level_sequence object
    binding = level_sequence.add_possessable(actor_to_animate)
    if not binding:
        # If it already exists, find it. This makes the script re-runnable.
        binding = level_sequence.find_possessable(actor_to_animate)

    # -----------------------------------------------------------------------------
    # 2. 获取或创建变换轨道
    # -----------------------------------------------------------------------------
    transform_track = None
    # Check existing tracks on the binding first
    for track in binding.get_tracks():
        if isinstance(track, unreal.MovieScene3DTransformTrack):
            transform_track = track
            break
    if not transform_track:
        transform_track = binding.add_track(unreal.MovieScene3DTransformTrack)

    # -----------------------------------------------------------------------------
    # 3. 设置轨道、范围和关键帧
    # -----------------------------------------------------------------------------
    # 清空旧数据，防止重复运行脚本时数据叠加
    for section in transform_track.get_sections():
        transform_track.remove_section(section)
    transform_section = transform_track.add_section()
    
    # MODIFIED: The frame rate and playback range are now set in main()
    # We only need to set the range for this specific actor's data section
    frame_rate = level_sequence.get_display_rate()
    fps = frame_rate.numerator / frame_rate.denominator
    start_frame = int(round(t[0] * fps))
    end_frame = int(round(t[-1] * fps))
    transform_section.set_range(start_frame, end_frame)

    # 获取位置和旋转的通道
    channels = transform_section.get_all_channels()
    location_x_channel, location_y_channel, location_z_channel = channels[0], channels[1], channels[2]
    rotation_x_channel, rotation_y_channel, rotation_z_channel = channels[3], channels[4], channels[5]

    unreal.log(f"开始为 {act_name} 添加关键帧...")
    # 【修正】使用 ScopedEditorTransaction 来创建可撤销的编辑器操作
    with unreal.ScopedEditorTransaction("Add Keys to Airplane Sequence") as trans:
        for i in range(len(t)):
            current_time = t[i]
            # MODIFIED: Apply UE5 scaling factor here
            current_pos = (X[i] * UE5_SCALING, Y[i] * UE5_SCALING, Z[i] * UE5_SCALING)
            current_rot = (Roll[i], Pitch[i], Yaw[i])

            frame_number = int(round(current_time * fps))
            frame_value = unreal.FrameNumber(frame_number)

            location_x_channel.add_key(frame_value, current_pos[0])
            location_y_channel.add_key(frame_value, current_pos[1])
            location_z_channel.add_key(frame_value, current_pos[2])

            rotation_x_channel.add_key(frame_value, current_rot[0]) # Roll
            rotation_y_channel.add_key(frame_value, current_rot[1]) # Pitch
            rotation_z_channel.add_key(frame_value, current_rot[2]) # Yaw

    unreal.log(f"{act_name} 的关键帧添加完成！")


def main():
    (
        all_time,
        all_longitudes,
        all_latitudes,
        all_altitudes,
        all_rolls,
        all_pitches,
        all_yaws
    ) = load_DCS_data()
    center, X, Y = global_to_local(
        all_longitudes,
        all_latitudes
    )

    # -----------------------------------------------------------------------------
    # NEW: Create or load the single master Level Sequence ONCE
    # -----------------------------------------------------------------------------
    asset_tools = unreal.AssetToolsHelpers.get_asset_tools()
    level_sequence = unreal.load_asset(SEQ_DIR + MASTER_SEQUENCE_NAME)
    if not level_sequence:
        level_sequence = asset_tools.create_asset(
            asset_name=MASTER_SEQUENCE_NAME,
            package_path=SEQ_DIR,
            asset_class=unreal.LevelSequence,
            factory=unreal.LevelSequenceFactoryNew()
        )
        unreal.log(f"已创建新的主关卡序列: {SEQ_DIR}{MASTER_SEQUENCE_NAME}")
    else:
        unreal.log(f"已加载现有的主关卡序列: {SEQ_DIR}{MASTER_SEQUENCE_NAME}")
        # Clear all old bindings to ensure a clean slate when re-running
        for binding in reversed(level_sequence.get_bindings()):
            level_sequence.remove_binding(binding)


    # -----------------------------------------------------------------------------
    # NEW: Calculate global time range and set up the sequence
    # -----------------------------------------------------------------------------
    frame_rate = unreal.FrameRate(30, 1) # 假设30fps
    level_sequence.set_display_rate(frame_rate)

    # Find the earliest start time and latest end time across all datasets
    global_start_time = float('inf')
    global_end_time = float('-inf')
    for plane_idx in all_time.keys():
        if all_time[plane_idx][0] < global_start_time:
            global_start_time = all_time[plane_idx][0]
        if all_time[plane_idx][-1] > global_end_time:
            global_end_time = all_time[plane_idx][-1]
    
    level_sequence.set_playback_start(global_start_time)
    level_sequence.set_playback_end(global_end_time)
    unreal.log(f"序列总播放范围已设置为: {global_start_time}s - {global_end_time}s")


    # -----------------------------------------------------------------------------
    # MODIFIED: Loop through planes and add them to the SAME sequence
    # -----------------------------------------------------------------------------
    for plane_idx in NAME_DATA_MAP.keys():
        add_actor_to_sequence(
            level_sequence, # Pass the master sequence
            plane_idx,
            all_time[plane_idx],
            X[plane_idx],
            Y[plane_idx],
            all_altitudes[plane_idx],
            all_rolls[plane_idx],
            all_pitches[plane_idx],
            all_yaws[plane_idx]
        )

    # -----------------------------------------------------------------------------
    # MOVED: Refresh the Sequencer window AFTER all actors have been added
    # -----------------------------------------------------------------------------
    unreal.get_editor_subsystem(unreal.AssetEditorSubsystem).open_editor_for_assets([level_sequence])
    unreal.log("脚本执行完毕，所有Actor已添加至主序列，并已打开或刷新Sequencer窗口。")


if __name__ == "__main__":
    main()
